from keras.models import Model
from keras.utils import plot_model
from keras.callbacks import EarlyStopping
from keras.layers import LSTM, Activation, Dense, Dropout, Input, Embedding
from keras.optimizers import RMSprop
from sklearn import metrics
import numpy as np
import sys
## 设置字体
from matplotlib.font_manager import FontProperties
fonts = FontProperties(fname = "/Library/Fonts/华文细黑.ttf",size=14)
import matplotlib.pyplot as plt
import seaborn as sns

class MyLSTM:
    def __init__(self,emb_size,vocab_size=5000,max_len=600,summary=True):
        self.vocab_size = vocab_size
        self.emb_size = emb_size
        self.max_len= max_len
        self.summary = summary

    def build_model(self):
        inputs = Input(name="inputs",shape=[self.max_len])
        # TODO Embedding(词汇表大小,batch大小,每个新闻的词长)
        emb = Embedding(input_dim=self.vocab_size + 1,output_dim= 128, input_length=self.max_len)(inputs)
        lstm = LSTM(units=128,name="LSTM_1")(emb)
        dense_1 = Dense(units=128,activation='relu',name="FC1")(lstm)
        dropout = Dropout(rate=0.5)(dense_1)
        output = Dense(units=10,activation='softmax',name="Classifier")(dropout)

        model = Model(inputs=inputs,outputs=output)
        if self.summary:
            model.summary()
        model.compile(loss="categorical_crossentropy", optimizer=RMSprop(), metrics=["accuracy"])
        plot_model(model, to_file="aa.png", show_shapes=True,dpi=300)
        return model

    def train(self,dataset,batch_size=128,epochs=10):

        train_seq_mat, train_y, val_seq_mat,val_y,test_seq_mat,test_y =dataset
        early_stop = EarlyStopping(monitor='val_loss',min_delta=0.0001)

        model = self.build_model()

        sys.exit(0)
        model_fit =model.fit(
            x=train_seq_mat,
            y=train_y,
            verbose=2,
            batch_size=128,
            epochs=10,
            validation_data=(val_seq_mat,val_y),
            callbacks=[early_stop]
        )
        self.test(model,test_seq_mat,test_y)
        model.save('my_model.h5')

    def test(self,model,test_x,test_y):
        test_pre = model.predict(test_x)
        ## 评价预测效果，计算混淆矩阵
        confm = metrics.confusion_matrix(np.argmax(test_pre, axis=1), np.argmax(test_y, axis=1))
        ## 混淆矩阵可视化
        Labname = ["体育", "娱乐", "家居", "房产", "教育", "时尚", "时政", "游戏", "科技", "财经"]
        plt.figure(figsize=(8, 8))
        sns.heatmap(confm.T, square=True, annot=True,
                    fmt='d', cbar=False, linewidths=.8,
                    cmap="YlGnBu")
        plt.xlabel('True label', size=14)
        plt.ylabel('Predicted label', size=14)
        plt.xticks(np.arange(10) + 0.5, Labname, fontproperties=fonts, size=12)
        plt.yticks(np.arange(10) + 0.3, Labname, fontproperties=fonts, size=12)
        plt.show()
        print(metrics.classification_report(np.argmax(test_pre, axis=1), np.argmax(test_y, axis=1)))
if __name__ == '__main__':
    m_lstm = MyLSTM(emb_size=200)
    m_lstm.build_model()


